%% Fraunhofer Diffraction and the 2D Fourier Transform;  Solving the Heat/Diffusion Equation with Fourier Series

%% Fraunhofer Diffraction and the 2D Fourier Transform
%
% Reference:
%%
% <https://en.wikipedia.org/wiki/Fraunhofer_diffraction_%28mathematics%29>
%
% The complex amplitude of a diffracted wave far from an aperture can be
% calculated using the Fraunhofer diffraction integral:
%%
%
% $$U(x,y,z)=\int\int_{Aperture} A(x',y') e^{-i   \frac{2\pi}{\lambda z} (x'x+y'y)} dx'dy'$$
%
% $x',y'$ are the coordinates across the aperture, $A(x',y')$ is the complex amplitude across the
% aperture, and $x,y,z$ are the coordinates at the screen.  The complex amplitude
% at each position $U(x,y,z)$ can be seen as the 2D Fourier coefficient
% calculated for the frequency $f_x=\frac{x}{\lambda z},f_y=\frac{y}{\lambda z}$.

%% The 2D Fourier transform
% The two dimensional fourier transform is computed using 'fft2'.  Let's
% try a simple example to demonstrate the 2D FT. 

lambda=5;  %Try adjusting and look at the result. 
k=2*pi/lambda;
[X,Y]=meshgrid(1:200,1:200);
Wave=real(exp(i*k*X));
figure;imshow(Wave,[]);title('Wave')

F_Wave=fft2(Wave);
figure;imshow(abs(F_Wave),[]);title('2D FFT of Wave')
F_Wave_Shift=fftshift(F_Wave);
figure;imshow(abs(F_Wave_Shift),[]);title('2D FFT of Wave, DC at Center.')

%% The Circular Aperture
% Let's revisit diffraction from a circular aperture. 

Lambda=.633;        %Wavelength of laser (micron)
D=50;               %Diameter of aperture (micron)
Z_Meters=.1;        %Screen distance in meters
Z=Z_Meters*10^6;    %Screen distance in microns

MeshSpacing=1;      %Sampling across aperture (micron)
MeshSize=200;       %Size of Screen (micron)

%Calculate and show the Amplitude across the aperture.
[XGrid,YGrid]=meshgrid((-MeshSize/2:MeshSpacing:MeshSize/2),(-MeshSize/2:MeshSpacing:MeshSize/2));
R=sqrt(XGrid.^2+YGrid.^2);
A=R<=D/2;  %The Amplitude across aperture (plane wave has constant amplitude, phase)
figure;imshow(A);title('Amplitude Across Aperture')

%Do the the 2D fft and show the result.
U=fftshift(fft2(A));
I=abs(U).^2;
figure;imshow(I,[]);title('Intensity at Screen')
figure;imshow(log(I),[]);title('Log Intensity at Screen')
figure;plot(I(101,:));xlabel('Pixels');ylabel('I')

% Now let's figure out the spacing of the pixels in I.
% The spacing of frequencies in the Fourier transform is 1/L
SZ=size(A,1);
FrequencySampling=1/(SZ*MeshSpacing); %units of 1/microns
[FGridX,FGridY]=meshgrid((0:SZ-1),(0:SZ-1));
FGridX_Shift=fftshift(FGridX);
FGridX_Shift(FGridX_Shift>SZ/2)=FGridX_Shift(FGridX_Shift>SZ/2)-SZ;
FGridY_Shift=fftshift(FGridY);
FGridY_Shift(FGridY_Shift>SZ/2)=FGridY_Shift(FGridY_Shift>SZ/2)-SZ;

FX=FGridX_Shift*FrequencySampling;
FY=FGridY_Shift*FrequencySampling;
X=FX*Lambda*Z;
Y=FY*Lambda*Z;

% Plot with units
I_norm=I/max(I(:));
X_center=X(101,:)+eps;
figure
plot(X_center,I_norm(101,:));xlabel('X (microns)');ylabel('I(X)')

% Compare to Airy disc
k=2*pi/Lambda;
a=D/2;
SinTheta=X_center./sqrt(X_center.^2+Z.^2);
I_Airy=(besselj(1,k*a*SinTheta)./(k*a*SinTheta)).^2;
I_Airy_norm=I_Airy/max(I_Airy(:));

hold on
plot(X_center,I_Airy_norm,'r--')
legend('Fourier','Direct')

%% The Finite-Width Single Slit
%

lambda=.633;        %Wavelength of laser (micron)
a=50;               %Diameter of aperture (micron)
Z_Meters=.1;        %Screen distance in meters
Z=Z_Meters*10^6;    %Screen distance in microns

MeshSpacing=1;  %Sampling across aperture (micron)
MeshSize=200;
[XGrid,YGrid]=meshgrid((-MeshSize/2:MeshSpacing:MeshSize/2),(-MeshSize/2:MeshSpacing:MeshSize/2));
R=abs(XGrid);
A=R<=(a/2);  %The Amplitude across aperture
figure;imshow(A);title('Amplitude Across Aperture')

U=fftshift(fft2(A));
I=abs(U).^2;
figure;imshow(I,[]);title('Intensity at Screen')
figure;imshow(log(I),[]);title('Log Intensity at Screen')
figure;plot(I(101,:));xlabel('Pixels');ylabel('I')

% Now let's figure out the spacing of the pixels in I.
% The spacing of frequencies in the Fourier transform is 1/L
SZ=size(A,1);
FrequencySampling=1/(SZ*MeshSpacing); %units of 1/microns
[FGridX,FGridY]=meshgrid((0:SZ-1),(0:SZ-1));
FGridX_Shift=fftshift(FGridX);
FGridX_Shift(FGridX_Shift>SZ/2)=FGridX_Shift(FGridX_Shift>SZ/2)-SZ;
FGridY_Shift=fftshift(FGridY);
FGridY_Shift(FGridY_Shift>SZ/2)=FGridY_Shift(FGridY_Shift>SZ/2)-SZ;

FX=FGridX_Shift*FrequencySampling;
FY=FGridY_Shift*FrequencySampling;
X=FX*Lambda*Z;
Y=FY*Lambda*Z;

% Plot with units
I_norm=I/max(I(:));
X_center=X(101,:)+eps;
figure
plot(X_center,I_norm(101,:));xlabel('X (microns)');ylabel('I(X)')

% Compare to Theory
k=2*pi/Lambda;
SinTheta=X_center./sqrt(X_center.^2+Z.^2);

%Theory (See also Mar_23.m)
% The FT of the 'rect' function is the 'sinc' function. 
%Make our own sinc function
mysinc=@(x)sin(pi*x)./(pi*x);
SinTheta=X_center./sqrt(Z.^2+X_center.^2);
SingleSlitDif=mysinc(SinTheta/Lambda*a).^2;

hold on
plot(X_center,SingleSlitDif,'r--')
legend('Fourier','Direct')

%% PDE Example: Diffusion in a box (Same as heat)
% We look the time-evolution of the distribution of small particles as they
% spread due to diffusion, but are confined to a 1D box. The box has length
% L. 
%
% The Diffusion equation is:
%%
% $$\frac{\partial u}{\partial t}=D\frac{\partial^2 u}{\partial x^2}$$
%
%
% Boundary conditions for the close box case are:
%
% $$\frac{\partial u}{\partial t}(0,t=0)$$,
% $$\frac{\partial u}{\partial t}(L,t=0)$$
%
% We can see that the following is a solution to the differential equation:
%
% $$u(x,t)=A_n e^{-\frac{n^2 D \pi^2}{L}}\cos{\frac{n pi x}{L}} $$
%
% where $A_n$ are determined by the intial distribution $u(x,0)$.

%Setup problem and initial state
L=100;          %Length of box
D=1;            %Diffusion Constant
U0=zeros(1,L);  %This is our t=0 particle distribution u(x,0)
U0(50)=1;       %Make this a delta function (we dropped a small drop of ink into the middle of 1D glass of water)

X=(0:L-1);
N=length(X)

%Calculate the Coeffiencients. 
A0=1/L*sum(U0);
for nn=1:N-1
    An(nn)=2/L*sum(U0.*cos(nn*pi*X/L));
end
An(N)=1/L*sum(U0.*cos(N*pi*X/L));
%Note the special end cases.  This correction is required to to exactly represent the U0 using an
%orthognal set of cosine base functions. The term An(N) is different in the
%discrete sampling case. 

% Calculates a term in the series sum. 
DiffEq = @(A,t,n,x)A*exp(-(n^2*D*pi^2*t/L^2))*cos(n*pi*x/L);

%Show the series sum to represent U0
X=(0:L-1);
t=0;
U=A0;
figure
for nn=1:N
    Wave=DiffEq(An(nn),t,nn,X);
    U=U+DiffEq(An(nn),t,nn,X);
    subplot(2,1,1);
    plot(X,Wave)
    xlabel('X');ylabel('Wave')
    subplot(2,1,2);
    hold off
    plot(X,U)
    hold on;
    plot(X,U0,'r--');
    xlabel('X');ylabel('U')
    %pause(.01)
end
legend('Fourier Sum','U0')

%% Show the time evolution:
% We simply evaluate the series at different time points. 

figure
for tt=0:.1:10
    U=A0;
    for nn=1:N
        U=U+DiffEq(An(nn),tt,nn,X);
    end
    plot(X,U) 
    xlabel('X');ylabel('U(x,t)')
    s=sprintf('time: %g',tt);
    text(80,.5,s);
    axis([0 L 0 1])
    %pause(1)
end
























